from typing import TypedDict, Annotated
import json
from langchain_community.tools import TavilySearchResults
from langchain_core.messages import ToolMessage
from langgraph.constants import START, END
from langgraph.graph.message import add_messages
from langgraph.graph.state import StateGraph

from langChain.config import model

tool = TavilySearchResults(max_results=2)
tools = [tool]
model_with_tools = model.bind_tools(tools)


# 要传递的状态
class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # 在注释中定义了该状态键应如何更新。
    # （在这种情况下，它将消息附加到列表中，而不是覆盖它们）
    messages: Annotated[list, add_messages]


graph_builder = StateGraph(State)


def chatbot(state: State):
    return {"messages": [model_with_tools.invoke(state["messages"])]}


# 第一个参数是唯一的节点名称。
# 第二个参数是将在每次调用时使用的函数或对象。
# 节点正在被使用。
graph_builder.add_node("chatbot", chatbot)


class BasicToolNode:
    """一个运行上一个AI消息中请求的工具的节点。"""

    def __init__(self, tools: list) -> None:
        self.tools_by_name = {tool.name: tool for tool in tools}

    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        outputs = []
        for tool_call in message.tool_calls:
            tool_result = self.tools_by_name[tool_call["name"]].invoke(
                tool_call["args"]
            )
            outputs.append(
                ToolMessage(
                    content=json.dumps(tool_result),
                    name=tool_call["name"],
                    tool_call_id=tool_call["id"],
                )
            )
        return {"messages": outputs}


tool_node = BasicToolNode(tools=[tool])
graph_builder.add_node("tools", tool_node)

graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)

graph = graph_builder.compile()

from IPython.display import Image, display


# try:
#     display(Image(graph.get_graph().draw_mermaid_png()))
# except Exception:
#     # 这需要一些额外的依赖，并且是可选的。
#     pass

def stream_graph_updates(user_input: str):
    for event in graph.stream({"messages": [("user", user_input)]}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)


while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        stream_graph_updates(user_input)
    except:
        # 如果 input() 不可用，则备选方案。
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break
